import qrcode
import io
import base64
from fastapi import Response
from fastapi.responses import JSONResponse

class QRCodeAPI:
    def __init__(self):
        pass
    
    def handle(self, text, box_size=10, border=4, fill_color="black", back_color="white", format="png"):
        """
        生成二维码并返回适当的响应
        
        参数:
        - text: 要编码的文本
        - box_size: 二维码中每个格子的像素大小
        - border: 二维码边框宽度(格子数)
        - fill_color: 二维码颜色
        - back_color: 背景颜色
        - format: 返回格式，'png'返回图片，'base64'返回base64编码
        
        返回:
        - Response对象或JSON响应
        """
        # 创建二维码
        qr = qrcode.QRCode(
            version=1,
            error_correction=qrcode.constants.ERROR_CORRECT_L,
            box_size=box_size,
            border=border,
        )
        qr.add_data(text)
        qr.make(fit=True)

        img = qr.make_image(fill_color=fill_color, back_color=back_color)
        
        # 根据format参数返回不同格式
        if format.lower() == "json":
            # 返回base64编码
            buffered = io.BytesIO()
            img.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
            return JSONResponse(content={"base64": f"data:image/png;base64,{img_str}"})
        elif format.lower() == "image":
            # 返回图片
            img_byte_arr = io.BytesIO()
            img.save(img_byte_arr, format='PNG')
            img_byte_arr.seek(0)
            
            return Response(content=img_byte_arr.getvalue(), media_type="image/png")